/*
 * UserNet  --  TCP/IP implementation in userspace
 * Copyright (C) 2010  Robert Millan
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdlib.h>		/* malloc */
#include <string.h>		/* memcpy */

#include <netinet/ip.h>		/* struct iphdr */
#include <netinet/ip_icmp.h>	/* struct icmphdr */
#include <netinet/tcp.h>	/* struct tcphdr */
#include "unet/checksum.h"

uint16_t
checksum (uint16_t *header, size_t len)
{
  uint32_t sum = 0;
  int i;

  for (i = 0; i < len / sizeof (uint16_t); i++)
    sum += ntohs (header[i]);

  return htons (~((sum >> 16) + (sum & 0xffff)));
}

int
unet_checksum_verify (void *packet)
{
  struct iphdr *header = packet;

  uint16_t packet_checksum = header->check;
  header->check = 0;
  uint16_t my_checksum = checksum ((void *) header, sizeof (struct iphdr));
  header->check = packet_checksum;

  return (packet_checksum == my_checksum);
}

static void
icmp_checksum_set (struct icmphdr *header)
{
  header->checksum = 0;
  header->checksum = checksum ((void *) header, sizeof (struct icmphdr));
}

struct tcp_ipv4_pseudo_header
{
  uint32_t saddr;
  uint32_t daddr;
  uint8_t zero;
  uint8_t protocol;
  uint16_t tcp_len;
} __attribute__ ((packed));

static void
tcp_checksum_set (uint8_t *ip_packet)
{
  struct iphdr *ip_header = (void *) ip_packet;
  uint16_t ip_len = ntohs (ip_header->tot_len);
  uint16_t tcp_len = ip_len - (ip_header->ihl * sizeof (uint32_t));
  uint8_t *tcp_packet = ip_packet + (ip_header->ihl * sizeof (uint32_t));
  struct tcphdr *tcp_header = (void *) tcp_packet;
  uint16_t ip_pseudo_len = tcp_len + sizeof (struct tcp_ipv4_pseudo_header);
  uint8_t *ip_pseudo_packet = malloc (ip_pseudo_len);
  struct tcp_ipv4_pseudo_header *ip_pseudo_header = (void *) ip_pseudo_packet;
  uint8_t *tcp_pseudo_packet = ip_pseudo_packet + sizeof (struct tcp_ipv4_pseudo_header);
  struct tcphdr *tcp_pseudo_header = (void *) tcp_pseudo_packet;

  ip_pseudo_header->saddr = ip_header->saddr;
  ip_pseudo_header->daddr = ip_header->daddr;
  ip_pseudo_header->zero = 0;
  ip_pseudo_header->protocol = IPPROTO_TCP;
  ip_pseudo_header->tcp_len = htons (tcp_len);

  memcpy (tcp_pseudo_packet, tcp_packet, tcp_len);

  tcp_pseudo_header->check = 0;

  tcp_header->check = checksum ((void *) ip_pseudo_packet, ip_pseudo_len);

  free (ip_pseudo_packet);
}

void
unet_checksum_set (void *packet)
{
  struct iphdr *header = packet;
  void *payload = packet + sizeof (*header);

  switch (header->protocol)
    {
    case IPPROTO_ICMP:
      icmp_checksum_set (payload);
      break;
    case IPPROTO_TCP:
      tcp_checksum_set (packet);
      break;
    }

  header->check = 0;
  header->check = checksum ((void *) header, sizeof (struct iphdr));
}
